import pandas as pd
from sklearn.model_selection import train_test_split
import os


def split_data_to_train_test(original_file_path, test_size=0.2, random_state=None):
    """
    将数据集随机切分为训练集和测试集

    参数:
    original_file_path (str): 原始数据集CSV文件的路径
    test_size (float): 测试集比例，默认0.2
    random_state (int): 随机种子，默认None

    返回:
    tuple: (train_file_path, test_file_path) 训练集和测试集的文件路径
    """
    try:
        # 读取数据
        df = pd.read_csv(original_file_path)

        # 随机切分
        train_df, test_df = train_test_split(
            df, test_size=test_size, random_state=random_state, shuffle=True
        )

        # 构建输出文件路径
        directory = os.path.dirname(original_file_path) or "."
        filename = os.path.splitext(os.path.basename(original_file_path))[0]

        train_path = os.path.join(directory, f"{filename}_train.csv")
        test_path = os.path.join(directory, f"{filename}_test.csv")

        # 保存文件
        train_df.to_csv(train_path, index=False)
        test_df.to_csv(test_path, index=False)

        return train_path, test_path

    except Exception as e:
        print(f"错误: {e}")
        return None, None


if __name__ == "__main__":
    # 配置参数
    original_file = "spam_dataset.csv"
    test_size = 0.2
    random_state = 42

    # 执行切分
    train_path, test_path = split_data_to_train_test(
        original_file, test_size=test_size, random_state=random_state
    )

    if train_path and test_path:
        print(f"数据集切分成功：")
        print(f"训练集: {train_path}")
        print(f"测试集: {test_path}")
    else:
        print("数据集切分失败")
